{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "stft requires the return_complex parameter be given for real inputs, and will further require that return_complex=True in a future PyTorch release.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m signal \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandn(\u001b[38;5;241m3\u001b[39m,\u001b[38;5;241m32000\u001b[39m)\n\u001b[1;32m----> 2\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstft\u001b[49m\u001b[43m(\u001b[49m\u001b[43msignal\u001b[49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Beilong Tang\\AppData\\Local\\Programs\\Python\\Python39\\lib\\site-packages\\torch\\functional.py:641\u001b[0m, in \u001b[0;36mstft\u001b[1;34m(input, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, return_complex)\u001b[0m\n\u001b[0;32m    639\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mpad(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39mview(extended_shape), [pad, pad], pad_mode)\n\u001b[0;32m    640\u001b[0m     \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39msignal_dim:])\n\u001b[1;32m--> 641\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_VF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstft\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_fft\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhop_length\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwin_length\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwindow\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[attr-defined]\u001b[39;49;00m\n\u001b[0;32m    642\u001b[0m \u001b[43m                \u001b[49m\u001b[43mnormalized\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43monesided\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreturn_complex\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: stft requires the return_complex parameter be given for real inputs, and will further require that return_complex=True in a future PyTorch release."
     ]
    }
   ],
   "source": [
    "signal = torch.randn(3,32000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.2222)\n"
     ]
    }
   ],
   "source": [
    "window_func = torch.hann_window(128)\n",
    "## model from their code\n",
    "output_old = torch.stft(signal,n_fft=128,win_length=128, hop_length=64,window=window_func,center=True,normalized=False,onesided=True,return_complex=True)\n",
    "output = torch.view_as_real(output_old)\n",
    "print(window_func[20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0590)\n",
      "tensor(0.2430)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 129, 501])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "window_func = torch.hann_window(256)\n",
    "print(window_func[20])\n",
    "\n",
    "window_func = torch.sqrt(torch.hann_window(256))\n",
    "print(window_func[20])\n",
    "# model from paper\n",
    "torch.stft(signal,n_fft=256,win_length=256, hop_length=64,window=window_func,center=True,normalized=False,onesided=True,return_complex=True).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 32, 501, 129])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(3,32,501,129)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,\n",
      "            11.,  12.,  13.,  14.],\n",
      "          [ 15.,  16.,  17.,  18.,  19.,  20.,  21.,  22.,  23.,  24.,  25.,\n",
      "            26.,  27.,  28.,  29.],\n",
      "          [ 30.,  31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,\n",
      "            41.,  42.,  43.,  44.],\n",
      "          [ 45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,  54.,  55.,\n",
      "            56.,  57.,  58.,  59.],\n",
      "          [ 60.,  61.,  62.,  63.,  64.,  65.,  66.,  67.,  68.,  69.,  70.,\n",
      "            71.,  72.,  73.,  74.],\n",
      "          [ 75.,  76.,  77.,  78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,\n",
      "            86.,  87.,  88.,  89.],\n",
      "          [ 90.,  91.,  92.,  93.,  94.,  95.,  96.,  97.,  98.,  99., 100.,\n",
      "           101., 102., 103., 104.],\n",
      "          [105., 106., 107., 108., 109., 110., 111., 112., 113., 114., 115.,\n",
      "           116., 117., 118., 119.],\n",
      "          [120., 121., 122., 123., 124., 125., 126., 127., 128., 129., 130.,\n",
      "           131., 132., 133., 134.],\n",
      "          [135., 136., 137., 138., 139., 140., 141., 142., 143., 144., 145.,\n",
      "           146., 147., 148., 149.],\n",
      "          [150., 151., 152., 153., 154., 155., 156., 157., 158., 159., 160.,\n",
      "           161., 162., 163., 164.],\n",
      "          [165., 166., 167., 168., 169., 170., 171., 172., 173., 174., 175.,\n",
      "           176., 177., 178., 179.],\n",
      "          [180., 181., 182., 183., 184., 185., 186., 187., 188., 189., 190.,\n",
      "           191., 192., 193., 194.],\n",
      "          [195., 196., 197., 198., 199., 200., 201., 202., 203., 204., 205.,\n",
      "           206., 207., 208., 209.],\n",
      "          [210., 211., 212., 213., 214., 215., 216., 217., 218., 219., 220.,\n",
      "           221., 222., 223., 224.]],\n",
      "\n",
      "         [[225., 226., 227., 228., 229., 230., 231., 232., 233., 234., 235.,\n",
      "           236., 237., 238., 239.],\n",
      "          [240., 241., 242., 243., 244., 245., 246., 247., 248., 249., 250.,\n",
      "           251., 252., 253., 254.],\n",
      "          [255., 256., 257., 258., 259., 260., 261., 262., 263., 264., 265.,\n",
      "           266., 267., 268., 269.],\n",
      "          [270., 271., 272., 273., 274., 275., 276., 277., 278., 279., 280.,\n",
      "           281., 282., 283., 284.],\n",
      "          [285., 286., 287., 288., 289., 290., 291., 292., 293., 294., 295.,\n",
      "           296., 297., 298., 299.],\n",
      "          [300., 301., 302., 303., 304., 305., 306., 307., 308., 309., 310.,\n",
      "           311., 312., 313., 314.],\n",
      "          [315., 316., 317., 318., 319., 320., 321., 322., 323., 324., 325.,\n",
      "           326., 327., 328., 329.],\n",
      "          [330., 331., 332., 333., 334., 335., 336., 337., 338., 339., 340.,\n",
      "           341., 342., 343., 344.],\n",
      "          [345., 346., 347., 348., 349., 350., 351., 352., 353., 354., 355.,\n",
      "           356., 357., 358., 359.],\n",
      "          [360., 361., 362., 363., 364., 365., 366., 367., 368., 369., 370.,\n",
      "           371., 372., 373., 374.],\n",
      "          [375., 376., 377., 378., 379., 380., 381., 382., 383., 384., 385.,\n",
      "           386., 387., 388., 389.],\n",
      "          [390., 391., 392., 393., 394., 395., 396., 397., 398., 399., 400.,\n",
      "           401., 402., 403., 404.],\n",
      "          [405., 406., 407., 408., 409., 410., 411., 412., 413., 414., 415.,\n",
      "           416., 417., 418., 419.],\n",
      "          [420., 421., 422., 423., 424., 425., 426., 427., 428., 429., 430.,\n",
      "           431., 432., 433., 434.],\n",
      "          [435., 436., 437., 438., 439., 440., 441., 442., 443., 444., 445.,\n",
      "           446., 447., 448., 449.]],\n",
      "\n",
      "         [[450., 451., 452., 453., 454., 455., 456., 457., 458., 459., 460.,\n",
      "           461., 462., 463., 464.],\n",
      "          [465., 466., 467., 468., 469., 470., 471., 472., 473., 474., 475.,\n",
      "           476., 477., 478., 479.],\n",
      "          [480., 481., 482., 483., 484., 485., 486., 487., 488., 489., 490.,\n",
      "           491., 492., 493., 494.],\n",
      "          [495., 496., 497., 498., 499., 500., 501., 502., 503., 504., 505.,\n",
      "           506., 507., 508., 509.],\n",
      "          [510., 511., 512., 513., 514., 515., 516., 517., 518., 519., 520.,\n",
      "           521., 522., 523., 524.],\n",
      "          [525., 526., 527., 528., 529., 530., 531., 532., 533., 534., 535.,\n",
      "           536., 537., 538., 539.],\n",
      "          [540., 541., 542., 543., 544., 545., 546., 547., 548., 549., 550.,\n",
      "           551., 552., 553., 554.],\n",
      "          [555., 556., 557., 558., 559., 560., 561., 562., 563., 564., 565.,\n",
      "           566., 567., 568., 569.],\n",
      "          [570., 571., 572., 573., 574., 575., 576., 577., 578., 579., 580.,\n",
      "           581., 582., 583., 584.],\n",
      "          [585., 586., 587., 588., 589., 590., 591., 592., 593., 594., 595.,\n",
      "           596., 597., 598., 599.],\n",
      "          [600., 601., 602., 603., 604., 605., 606., 607., 608., 609., 610.,\n",
      "           611., 612., 613., 614.],\n",
      "          [615., 616., 617., 618., 619., 620., 621., 622., 623., 624., 625.,\n",
      "           626., 627., 628., 629.],\n",
      "          [630., 631., 632., 633., 634., 635., 636., 637., 638., 639., 640.,\n",
      "           641., 642., 643., 644.],\n",
      "          [645., 646., 647., 648., 649., 650., 651., 652., 653., 654., 655.,\n",
      "           656., 657., 658., 659.],\n",
      "          [660., 661., 662., 663., 664., 665., 666., 667., 668., 669., 670.,\n",
      "           671., 672., 673., 674.]]]])\n",
      "torch.Size([1, 27, 169])\n",
      "tensor([[[[[  0.,   1.,   2.],\n",
      "           [ 15.,  16.,  17.],\n",
      "           [ 30.,  31.,  32.]],\n",
      "\n",
      "          [[225., 226., 227.],\n",
      "           [240., 241., 242.],\n",
      "           [255., 256., 257.]],\n",
      "\n",
      "          [[450., 451., 452.],\n",
      "           [465., 466., 467.],\n",
      "           [480., 481., 482.]]],\n",
      "\n",
      "\n",
      "         [[[  1.,   2.,   3.],\n",
      "           [ 16.,  17.,  18.],\n",
      "           [ 31.,  32.,  33.]],\n",
      "\n",
      "          [[226., 227., 228.],\n",
      "           [241., 242., 243.],\n",
      "           [256., 257., 258.]],\n",
      "\n",
      "          [[451., 452., 453.],\n",
      "           [466., 467., 468.],\n",
      "           [481., 482., 483.]]],\n",
      "\n",
      "\n",
      "         [[[  2.,   3.,   4.],\n",
      "           [ 17.,  18.,  19.],\n",
      "           [ 32.,  33.,  34.]],\n",
      "\n",
      "          [[227., 228., 229.],\n",
      "           [242., 243., 244.],\n",
      "           [257., 258., 259.]],\n",
      "\n",
      "          [[452., 453., 454.],\n",
      "           [467., 468., 469.],\n",
      "           [482., 483., 484.]]],\n",
      "\n",
      "\n",
      "         ...,\n",
      "\n",
      "\n",
      "         [[[190., 191., 192.],\n",
      "           [205., 206., 207.],\n",
      "           [220., 221., 222.]],\n",
      "\n",
      "          [[415., 416., 417.],\n",
      "           [430., 431., 432.],\n",
      "           [445., 446., 447.]],\n",
      "\n",
      "          [[640., 641., 642.],\n",
      "           [655., 656., 657.],\n",
      "           [670., 671., 672.]]],\n",
      "\n",
      "\n",
      "         [[[191., 192., 193.],\n",
      "           [206., 207., 208.],\n",
      "           [221., 222., 223.]],\n",
      "\n",
      "          [[416., 417., 418.],\n",
      "           [431., 432., 433.],\n",
      "           [446., 447., 448.]],\n",
      "\n",
      "          [[641., 642., 643.],\n",
      "           [656., 657., 658.],\n",
      "           [671., 672., 673.]]],\n",
      "\n",
      "\n",
      "         [[[192., 193., 194.],\n",
      "           [207., 208., 209.],\n",
      "           [222., 223., 224.]],\n",
      "\n",
      "          [[417., 418., 419.],\n",
      "           [432., 433., 434.],\n",
      "           [447., 448., 449.]],\n",
      "\n",
      "          [[642., 643., 644.],\n",
      "           [657., 658., 659.],\n",
      "           [672., 673., 674.]]]]])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from torch.nn import functional as f\n",
    " \n",
    "x = torch.arange(0, 1*3*15*15).float()\n",
    "x = x.view(1,3,15,15)\n",
    "print(x)\n",
    "x1 = f.unfold(x, kernel_size=3, dilation=1, stride=1)\n",
    "print(x1.shape)\n",
    "B, C_kh_kw, L = x1.size()\n",
    "x1 = x1.permute(0, 2, 1)\n",
    "x1 = x1.view(B, L, -1, 3, 3)\n",
    "print(x1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 800, 62125])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = torch.nn.functional.unfold(x,kernel_size=(5,5))\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(4.2488-0.0830j)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_old[0,55,100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 3.9868e+01,  0.0000e+00],\n",
       "          [ 5.9999e+00,  0.0000e+00],\n",
       "          [ 7.5186e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [-3.6830e+01,  0.0000e+00],\n",
       "          [-5.9023e+01,  0.0000e+00],\n",
       "          [-7.3112e+01,  0.0000e+00]],\n",
       "\n",
       "         [[ 4.0918e+01,  0.0000e+00],\n",
       "          [-1.8271e+01,  1.5838e+01],\n",
       "          [-7.8137e+00, -2.4140e+01],\n",
       "          ...,\n",
       "          [ 5.7981e+00, -1.7037e+01],\n",
       "          [ 4.7798e+00, -1.0717e+01],\n",
       "          [-3.7321e-01, -4.5786e-03]],\n",
       "\n",
       "         [[-1.1014e+01,  9.5367e-07],\n",
       "          [ 1.9737e+01, -2.2088e+01],\n",
       "          [-2.9122e+00,  2.3470e+01],\n",
       "          ...,\n",
       "          [ 2.0038e+01, -1.1870e+01],\n",
       "          [-1.4956e+01, -3.8427e+00],\n",
       "          [ 8.1432e+00,  1.9990e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.1493e+01,  0.0000e+00],\n",
       "          [-1.5087e+01, -1.1341e+01],\n",
       "          [ 1.1377e+01,  1.4683e+01],\n",
       "          ...,\n",
       "          [-3.3413e+00,  8.5074e+00],\n",
       "          [-2.4737e+00, -3.0121e+00],\n",
       "          [ 9.1654e+00, -2.2500e-01]],\n",
       "\n",
       "         [[-2.8458e+01,  0.0000e+00],\n",
       "          [ 4.1824e+00,  2.0977e+01],\n",
       "          [ 1.2327e+01,  8.5836e-01],\n",
       "          ...,\n",
       "          [ 5.1228e+00,  9.9297e+00],\n",
       "          [-2.5119e+00,  2.8950e+00],\n",
       "          [ 1.0850e+01, -1.3315e-01]],\n",
       "\n",
       "         [[-2.6411e+01,  0.0000e+00],\n",
       "          [-1.9537e+01,  0.0000e+00],\n",
       "          [-3.1637e+01,  0.0000e+00],\n",
       "          ...,\n",
       "          [-3.9403e+01,  0.0000e+00],\n",
       "          [-5.7467e+01,  0.0000e+00],\n",
       "          [-4.8240e+01,  0.0000e+00]]],\n",
       "\n",
       "\n",
       "        [[[-1.9836e+01,  0.0000e+00],\n",
       "          [-1.4625e+01,  0.0000e+00],\n",
       "          [-8.4066e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [-3.3208e+01,  0.0000e+00],\n",
       "          [-3.6489e+01,  0.0000e+00],\n",
       "          [-5.5513e+01,  0.0000e+00]],\n",
       "\n",
       "         [[ 8.2081e-01, -1.9073e-06],\n",
       "          [ 8.2095e+00, -1.5547e+00],\n",
       "          [ 7.4647e+00,  1.2079e+01],\n",
       "          ...,\n",
       "          [-3.5865e+00, -8.1264e+00],\n",
       "          [ 1.0244e+01, -7.4066e+00],\n",
       "          [-3.8602e+00, -4.7375e-02]],\n",
       "\n",
       "         [[ 2.7834e+01, -1.9073e-06],\n",
       "          [-1.6000e+01,  1.8239e+00],\n",
       "          [ 1.9970e+01,  6.9655e+00],\n",
       "          ...,\n",
       "          [ 2.4633e+01,  9.8691e+00],\n",
       "          [-1.9677e+01, -5.9425e+00],\n",
       "          [ 1.7540e+01,  4.3058e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-3.4684e+01,  3.8147e-06],\n",
       "          [ 1.7061e+01,  6.1012e-01],\n",
       "          [-9.3842e+00, -1.4769e+01],\n",
       "          ...,\n",
       "          [-1.6196e+01,  6.8442e+00],\n",
       "          [ 6.9856e+00, -1.6717e+01],\n",
       "          [-6.1390e+00,  1.5071e-01]],\n",
       "\n",
       "         [[ 1.9799e+01,  1.9073e-06],\n",
       "          [-8.1918e+00, -2.6246e+01],\n",
       "          [-4.5486e+00, -8.1981e+00],\n",
       "          ...,\n",
       "          [ 3.3750e+00, -1.4633e+01],\n",
       "          [-6.9161e+00, -1.9402e+01],\n",
       "          [-1.3161e+01,  1.6152e-01]],\n",
       "\n",
       "         [[-5.6928e+00,  0.0000e+00],\n",
       "          [-7.1016e+00,  0.0000e+00],\n",
       "          [ 2.4728e+01,  0.0000e+00],\n",
       "          ...,\n",
       "          [ 8.0907e+00,  0.0000e+00],\n",
       "          [ 2.7533e+01,  0.0000e+00],\n",
       "          [ 3.3683e+01,  0.0000e+00]]],\n",
       "\n",
       "\n",
       "        [[[ 5.6128e+01,  0.0000e+00],\n",
       "          [ 3.4400e+01,  0.0000e+00],\n",
       "          [ 4.2794e+00,  0.0000e+00],\n",
       "          ...,\n",
       "          [ 1.6845e+01,  0.0000e+00],\n",
       "          [ 1.6628e+01,  0.0000e+00],\n",
       "          [ 2.1028e+01,  0.0000e+00]],\n",
       "\n",
       "         [[-1.1965e+01, -2.3842e-07],\n",
       "          [-1.7785e+01, -2.3610e+01],\n",
       "          [ 2.3607e+00, -3.6005e+01],\n",
       "          ...,\n",
       "          [ 1.2066e+01, -9.4922e+00],\n",
       "          [ 5.8007e+00,  1.4028e+01],\n",
       "          [-3.5405e+00, -4.3451e-02]],\n",
       "\n",
       "         [[-1.5804e+01,  0.0000e+00],\n",
       "          [ 7.8589e+00, -1.8132e+01],\n",
       "          [-9.9883e+00, -3.8523e+00],\n",
       "          ...,\n",
       "          [-1.0243e+01, -3.4023e+00],\n",
       "          [ 5.4465e+00, -5.8233e-02],\n",
       "          [ 1.4258e+01,  3.5001e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.4614e+01,  0.0000e+00],\n",
       "          [-1.0376e+01, -1.0574e+01],\n",
       "          [ 2.2021e+01, -6.7542e-01],\n",
       "          ...,\n",
       "          [-1.1457e+01,  8.4614e+00],\n",
       "          [ 1.5953e+01, -1.2735e+01],\n",
       "          [-1.7287e+01,  4.2438e-01]],\n",
       "\n",
       "         [[-8.8132e+00, -9.5367e-07],\n",
       "          [ 9.1694e+00,  3.2357e+00],\n",
       "          [ 1.8907e+01, -1.5560e+01],\n",
       "          ...,\n",
       "          [ 9.2566e+00, -1.7665e+01],\n",
       "          [-1.4394e+01, -9.9955e+00],\n",
       "          [ 3.0010e+00, -3.6828e-02]],\n",
       "\n",
       "         [[-1.1345e+01,  0.0000e+00],\n",
       "          [-9.8551e-01,  0.0000e+00],\n",
       "          [ 1.6264e+01,  0.0000e+00],\n",
       "          ...,\n",
       "          [-1.8555e+01,  0.0000e+00],\n",
       "          [-1.6275e+01,  0.0000e+00],\n",
       "          [ 5.8031e+00,  0.0000e+00]]]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.stft(signal,512,return_complex=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = torch.randn(3,256,34,33)\n",
    "import torch.nn.functional as F\n",
    "\n",
    "output = F.pad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[...,None].shape\n",
    "x[...,None][0,0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 2, 25, 32])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn(3,2,25,32)\n",
    "\n",
    "layerNorm = torch.nn.LayerNorm([25,32])\n",
    "y = layerNorm(x)\n",
    "print(y.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 2, 4])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "x_test = np.array([[[1,2,-1,1],[3,4,-2,2]],\n",
    "                   [[1,2,-1,1],[3,4,-2,2]]])\n",
    "x_test = torch.from_numpy(x_test).float()\n",
    "x_test = x_test.unsqueeze(0)\n",
    "x_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.1348,  0.4045, -1.2136, -0.1348],\n",
       "          [ 0.9439,  1.4832, -1.7529,  0.4045]],\n",
       "\n",
       "         [[-0.1348,  0.4045, -1.2136, -0.1348],\n",
       "          [ 0.9439,  1.4832, -1.7529,  0.4045]]]],\n",
       "       grad_fn=<NativeLayerNormBackward0>)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = torch.nn.LayerNorm(normalized_shape = [2,4])\n",
    "output = m(x_test)\n",
    "output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.1348,  0.4045, -1.2136, -0.1348],\n",
       "          [ 0.9439,  1.4832, -1.7529,  0.4045]],\n",
       "\n",
       "         [[-0.1348,  0.4045, -1.2136, -0.1348],\n",
       "          [ 0.9439,  1.4832, -1.7529,  0.4045]]]],\n",
       "       grad_fn=<NativeGroupNormBackward0>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = torch.nn.GroupNorm(1,2)\n",
    "output = g(x_test)\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "class LayerNormalization4DCF(nn.Module):\n",
    "    def __init__(self, input_dimension, eps=1e-5):\n",
    "        super().__init__()\n",
    "        assert len(input_dimension) == 2\n",
    "        param_size = [1, input_dimension[0], 1, input_dimension[1]]\n",
    "        self.gamma = nn.parameter.Parameter(torch.Tensor(*param_size).to(torch.float32))\n",
    "        self.beta = nn.parameter.Parameter(torch.Tensor(*param_size).to(torch.float32))\n",
    "        nn.init.ones_(self.gamma)\n",
    "        nn.init.zeros_(self.beta)\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        if x.ndim == 4:\n",
    "            stat_dim = (1, 3)\n",
    "        else:\n",
    "            raise ValueError(\"Expect x to have 4 dimensions, but got {}\".format(x.ndim))\n",
    "        mu_ = x.mean(dim=stat_dim, keepdim=True)  # [B,1,T,1]\n",
    "        std_ = torch.sqrt(\n",
    "            x.var(dim=stat_dim, unbiased=False, keepdim=True) + self.eps\n",
    "        )  # [B,1,T,F]\n",
    "        x_hat = ((x - mu_) / std_) * self.gamma + self.beta\n",
    "        return x_hat\n",
    "x = torch.randn(3,4,501,129)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 9.1111e-01, -2.6965e+00, -4.5666e-01,  ...,  1.1969e+00,\n",
       "            2.6031e-01,  1.4139e+00],\n",
       "          [ 5.2745e-01, -1.1000e+00,  6.5941e-01,  ..., -7.1783e-01,\n",
       "           -7.3384e-02,  9.5959e-01],\n",
       "          [-6.0912e-01,  1.5149e+00, -2.1327e+00,  ..., -7.4061e-01,\n",
       "           -3.4561e-01, -1.2236e+00],\n",
       "          [-5.3389e-01, -1.4661e+00,  1.1110e+00,  ...,  9.7702e-01,\n",
       "           -1.2032e+00,  1.6727e+00]],\n",
       "\n",
       "         [[ 3.2614e-01, -1.3608e+00, -5.1389e-01,  ..., -2.0681e+00,\n",
       "            3.2315e-01,  1.8129e+00],\n",
       "          [-1.1976e+00,  8.2611e-01, -1.0249e+00,  ...,  6.5088e-01,\n",
       "            4.7706e-02,  1.3192e+00],\n",
       "          [-6.7538e-01, -4.7069e-01, -4.0662e-01,  ...,  2.8753e-02,\n",
       "           -2.2253e+00, -1.0983e+00],\n",
       "          [ 1.8557e+00, -4.6133e-01,  1.6033e+00,  ...,  4.7855e-01,\n",
       "            1.5378e+00,  1.1973e+00]],\n",
       "\n",
       "         [[ 2.3268e-01, -1.5906e+00,  2.4693e+00,  ...,  1.6851e-01,\n",
       "           -1.0819e-01,  3.3279e-01],\n",
       "          [ 3.2019e-01,  3.9563e-01, -6.3921e-01,  ..., -4.3906e-01,\n",
       "            6.6902e-01,  2.9696e-01],\n",
       "          [ 9.8413e-01,  4.5077e-01, -3.4549e-02,  ..., -7.5173e-01,\n",
       "            1.4281e+00,  1.2289e+00],\n",
       "          [-3.2796e-01, -1.3515e+00,  7.4165e-01,  ...,  1.6122e+00,\n",
       "           -1.4221e+00, -8.6541e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.4136e+00,  8.2688e-01,  7.8341e-01,  ..., -1.0988e+00,\n",
       "            9.8616e-01, -1.0983e+00],\n",
       "          [-2.3953e+00,  2.8159e+00, -1.1616e+00,  ..., -5.5926e-01,\n",
       "            2.3331e+00, -1.8304e+00],\n",
       "          [ 7.5514e-01,  1.0211e-01, -5.9772e-01,  ...,  7.1456e-01,\n",
       "           -4.5549e-01,  6.9860e-01],\n",
       "          [ 1.8110e-01,  1.8494e+00,  5.8072e-01,  ...,  5.3070e-02,\n",
       "            5.5005e-01, -7.2534e-02]],\n",
       "\n",
       "         [[-9.7365e-01,  1.0096e+00,  6.0439e-01,  ..., -2.5716e+00,\n",
       "           -5.9073e-01,  1.1478e+00],\n",
       "          [-2.3208e+00,  7.9556e-01,  7.2947e-01,  ...,  2.4091e+00,\n",
       "           -8.5328e-01,  8.8391e-01],\n",
       "          [ 8.8119e-01, -1.2612e+00, -1.8960e+00,  ..., -1.9696e-01,\n",
       "            1.4735e+00, -7.9435e-01],\n",
       "          [ 1.9057e+00, -5.1708e-01,  3.1024e-01,  ...,  2.3652e+00,\n",
       "            5.8461e-03, -9.8625e-01]],\n",
       "\n",
       "         [[ 1.3853e-01,  4.4441e-01,  9.5448e-01,  ..., -4.7053e-01,\n",
       "            1.2693e+00,  1.1187e+00],\n",
       "          [-4.2867e-01, -3.9381e-01,  6.8732e-01,  ...,  1.4462e-01,\n",
       "           -4.2888e-01, -8.9005e-01],\n",
       "          [-2.0830e+00, -7.3713e-01,  2.4213e+00,  ...,  1.2683e-01,\n",
       "           -5.1553e-01,  1.4088e-01],\n",
       "          [-2.9260e-01, -4.7313e-03, -2.5689e-01,  ...,  1.9638e-01,\n",
       "           -7.5804e-01,  9.5841e-01]]],\n",
       "\n",
       "\n",
       "        [[[ 4.3302e-01, -7.6274e-01,  3.9616e-01,  ...,  9.7908e-02,\n",
       "           -5.2624e-01,  9.1276e-01],\n",
       "          [-4.5934e-01, -1.1039e-01, -7.5764e-01,  ...,  5.1770e-01,\n",
       "            6.4844e-01,  1.2536e-02],\n",
       "          [ 5.8030e-01,  4.5078e-01,  9.5006e-03,  ..., -8.1235e-01,\n",
       "            1.2955e-01,  1.0350e+00],\n",
       "          [-9.2550e-01,  2.6215e-01, -1.4791e-02,  ..., -5.6148e-01,\n",
       "            7.3604e-01, -5.5321e-01]],\n",
       "\n",
       "         [[-9.5820e-01,  8.5652e-01, -2.9300e-01,  ..., -1.1565e+00,\n",
       "           -3.5107e-01,  6.0001e-01],\n",
       "          [-9.7830e-01, -4.6672e-01,  9.6535e-02,  ...,  1.2761e+00,\n",
       "           -2.1593e+00, -1.2637e-01],\n",
       "          [ 1.5451e+00, -8.0602e-02, -2.5314e-03,  ..., -2.5113e-01,\n",
       "           -1.6936e+00,  9.2470e-01],\n",
       "          [-1.3478e-01,  1.1366e+00,  4.3186e-01,  ...,  4.3627e-01,\n",
       "            7.2640e-01, -4.2234e-02]],\n",
       "\n",
       "         [[ 4.2419e-01, -6.8301e-01,  1.1143e+00,  ..., -6.2594e-01,\n",
       "           -2.1925e+00, -9.1203e-01],\n",
       "          [ 5.1888e-01, -3.9706e-01,  2.9723e-01,  ..., -2.1562e+00,\n",
       "           -8.2261e-01,  2.2414e+00],\n",
       "          [-2.1849e+00,  1.0960e-01,  7.0246e-01,  ..., -1.2312e+00,\n",
       "            7.0479e-01, -8.8975e-01],\n",
       "          [ 1.1756e+00, -7.9862e-01,  5.4169e-01,  ..., -4.3449e-01,\n",
       "            1.2133e-02, -7.8238e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-8.9124e-02,  1.4588e-01, -5.3723e-01,  ...,  3.3059e-01,\n",
       "           -1.9651e+00, -1.6192e+00],\n",
       "          [-1.0115e-01, -3.3749e-01,  1.3768e+00,  ..., -1.3614e-01,\n",
       "            3.8908e-01,  1.4969e-01],\n",
       "          [-1.3222e+00,  1.6016e+00,  2.5915e-01,  ..., -9.2491e-01,\n",
       "            5.7427e-01, -5.3615e-01],\n",
       "          [ 9.9903e-01, -2.3357e-01,  1.9838e+00,  ..., -4.7926e-01,\n",
       "           -4.8890e-01,  1.4789e+00]],\n",
       "\n",
       "         [[-1.1650e+00,  1.9409e+00, -2.8858e+00,  ...,  1.6811e+00,\n",
       "           -1.1902e-01, -5.2911e-01],\n",
       "          [ 4.3555e-01, -2.8404e-01,  1.5505e-01,  ..., -1.0329e+00,\n",
       "           -1.5034e+00, -2.6291e+00],\n",
       "          [ 9.1508e-01, -1.7520e-01, -8.9358e-01,  ...,  5.0356e-01,\n",
       "            1.6227e+00, -5.7818e-01],\n",
       "          [-7.3089e-01,  1.1442e+00,  5.1752e-01,  ...,  6.0569e-01,\n",
       "           -5.5498e-01, -2.3249e+00]],\n",
       "\n",
       "         [[ 5.0750e-01, -1.5561e+00, -1.6975e-02,  ...,  8.1543e-01,\n",
       "           -5.6799e-01, -9.2776e-01],\n",
       "          [-4.1853e-01,  4.1284e-01, -6.1291e-02,  ..., -3.2761e-01,\n",
       "           -5.5534e-01,  1.1326e+00],\n",
       "          [-8.1934e-01, -2.2986e-01, -5.6026e-01,  ..., -1.8076e+00,\n",
       "            1.9420e+00,  1.5542e+00],\n",
       "          [ 1.3434e+00,  1.7872e-01, -9.9169e-01,  ..., -1.7403e+00,\n",
       "           -4.3832e-01, -1.5404e+00]]],\n",
       "\n",
       "\n",
       "        [[[-1.2347e+00, -1.8745e-01,  1.0209e-01,  ...,  1.1629e+00,\n",
       "            3.0225e-01, -7.2551e-01],\n",
       "          [-6.0164e-02, -2.2165e-01,  4.2509e-01,  ..., -6.4197e-02,\n",
       "           -1.8118e-01,  2.5436e-01],\n",
       "          [ 3.0178e-01,  4.4534e-01, -2.7927e-01,  ...,  1.6257e+00,\n",
       "           -5.1702e-01,  1.4840e+00],\n",
       "          [ 1.2067e-01,  2.3424e-01, -2.9598e-01,  ...,  1.3570e+00,\n",
       "            3.2527e-01, -2.2949e+00]],\n",
       "\n",
       "         [[ 4.0171e-01,  1.5921e+00, -2.6131e+00,  ...,  2.9751e-01,\n",
       "            1.7066e+00, -2.5286e-01],\n",
       "          [-5.1508e-01, -9.9591e-01,  1.9569e-01,  ...,  1.0416e-01,\n",
       "           -6.8788e-01, -1.1415e-03],\n",
       "          [ 4.7295e-01,  1.0062e+00,  2.1192e+00,  ..., -8.4962e-01,\n",
       "            1.3210e+00, -2.8761e-01],\n",
       "          [ 4.4459e-02, -6.7586e-01,  2.2065e-01,  ...,  1.2022e+00,\n",
       "           -4.4560e-01, -9.5140e-01]],\n",
       "\n",
       "         [[-1.8881e+00,  7.9569e-01, -7.1445e-01,  ...,  6.0761e-01,\n",
       "           -1.4854e+00,  3.8005e-01],\n",
       "          [-6.8997e-01,  4.1320e-01, -5.8285e-01,  ...,  2.4037e-01,\n",
       "            2.1801e+00, -1.9739e+00],\n",
       "          [ 1.2673e+00,  3.2284e-01, -2.6383e-01,  ..., -1.7443e+00,\n",
       "           -4.9880e-01,  1.6476e+00],\n",
       "          [-4.2736e-02,  1.1388e-01,  1.0239e+00,  ...,  4.9253e-01,\n",
       "           -2.0428e+00,  7.2641e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 3.2008e-01, -2.2604e-01,  1.2194e-01,  ...,  5.0447e-01,\n",
       "           -1.4719e+00,  7.7064e-02],\n",
       "          [-3.0133e-01, -1.7949e+00,  6.0130e-01,  ...,  3.8042e-02,\n",
       "            1.1178e-01,  1.7755e-01],\n",
       "          [ 4.6276e-01,  7.8358e-01, -1.0370e+00,  ...,  3.5579e-02,\n",
       "           -1.9842e+00,  4.0497e-01],\n",
       "          [-3.5856e-01, -9.8944e-01,  2.1008e-01,  ...,  1.5221e+00,\n",
       "           -7.7752e-01,  2.1349e-01]],\n",
       "\n",
       "         [[ 2.3341e+00, -8.8860e-01, -1.3656e+00,  ...,  4.8946e-01,\n",
       "           -4.2845e-01,  9.2548e-01],\n",
       "          [ 2.9260e-01, -3.8432e-01,  2.8407e-01,  ..., -1.5215e+00,\n",
       "           -8.1998e-01,  1.2565e+00],\n",
       "          [ 2.0910e-01,  6.3109e-01,  1.2059e+00,  ...,  4.9602e-01,\n",
       "           -1.4742e+00, -6.7826e-01],\n",
       "          [ 3.4276e-02,  6.0279e-01,  2.0881e+00,  ...,  1.6545e+00,\n",
       "            1.4105e+00,  5.9151e-01]],\n",
       "\n",
       "         [[ 2.5955e-01, -2.4873e+00,  4.1470e-01,  ..., -1.4958e-01,\n",
       "            1.0128e+00, -5.5772e-01],\n",
       "          [ 5.0954e-02,  9.0936e-01, -6.3618e-01,  ..., -4.3758e-01,\n",
       "            2.6359e-01,  5.7093e-02],\n",
       "          [-6.3663e-01, -9.3113e-01,  1.6905e+00,  ..., -3.9296e-01,\n",
       "            2.8493e-01,  1.4122e+00],\n",
       "          [-8.9018e-01,  7.7206e-01, -8.1522e-01,  ...,  1.1613e+00,\n",
       "           -1.1636e+00,  1.4999e+00]]]], grad_fn=<TransposeBackward0>)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "layerNorm = LayerNormalization4DCF((4,129))\n",
    "y = layerNorm(x).transpose(1,2)\n",
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 501, 4, 129])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l = nn.LayerNorm((4,129))\n",
    "l(x.transpose(1,2).contiguous()).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = torch.randn(3, 501, 4128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 1.4173e+00,  7.7962e-01,  3.9546e-03,  ..., -2.7568e-02,\n",
       "           -9.6422e-01, -3.3995e-01],\n",
       "          [-4.1776e-01,  2.0515e-01, -3.2592e-01,  ..., -5.0554e-01,\n",
       "           -8.5626e-01, -7.5003e-01],\n",
       "          [-8.8587e-01,  1.3755e+00,  1.5670e+00,  ..., -6.8439e-01,\n",
       "           -8.3740e-01,  3.2749e-01],\n",
       "          ...,\n",
       "          [ 6.0865e-01, -8.3371e-01,  5.3483e-01,  ...,  4.8943e-01,\n",
       "            5.2437e-01, -8.0147e-01],\n",
       "          [-2.1721e+00,  5.0381e-01,  4.4588e-01,  ...,  1.3669e+00,\n",
       "            8.3749e-02, -1.1268e+00],\n",
       "          [-3.5146e-01, -2.8747e-01,  3.6681e-01,  ..., -6.7973e-01,\n",
       "           -1.1825e+00, -1.7310e+00]],\n",
       "\n",
       "         [[ 1.1254e+00, -7.5612e-02,  4.9629e-01,  ..., -5.2587e-01,\n",
       "            1.1161e+00,  1.6859e-01],\n",
       "          [-5.8612e-01,  1.1333e+00,  1.2494e+00,  ..., -3.2623e-01,\n",
       "           -3.9444e-01,  3.8408e-01],\n",
       "          [-2.0357e+00,  3.0590e-01, -2.2440e+00,  ...,  1.1050e+00,\n",
       "            1.9836e-01, -6.8045e-02],\n",
       "          ...,\n",
       "          [ 4.0967e-02,  2.3345e-01,  7.8230e-01,  ...,  8.4325e-02,\n",
       "            1.1821e+00, -3.8880e-01],\n",
       "          [ 2.7598e-01,  8.1544e-01, -4.1195e-01,  ...,  8.6941e-01,\n",
       "            8.7553e-02, -7.6473e-01],\n",
       "          [ 8.3201e-01,  2.1823e+00,  3.9377e-01,  ..., -6.5179e-01,\n",
       "            1.3267e+00,  1.3558e-01]],\n",
       "\n",
       "         [[ 8.3895e-01, -7.6690e-01,  4.9109e-01,  ...,  5.9498e-02,\n",
       "           -1.1891e+00,  2.7977e-01],\n",
       "          [ 4.8181e-01, -1.5749e+00, -4.3431e-01,  ..., -6.7664e-01,\n",
       "            3.3375e-01,  7.0330e-01],\n",
       "          [ 8.2954e-01,  3.5906e-01, -1.2993e-02,  ..., -3.0799e-01,\n",
       "            4.2813e-01, -1.0685e+00],\n",
       "          ...,\n",
       "          [ 1.2054e+00, -3.4316e-01, -6.4021e-01,  ..., -1.4281e+00,\n",
       "            2.3123e-01,  7.3876e-02],\n",
       "          [-1.3922e-01,  1.1440e+00,  9.4357e-01,  ...,  2.2402e-01,\n",
       "            1.1697e+00, -1.0227e+00],\n",
       "          [-6.9713e-03, -2.2048e-01, -3.2757e-01,  ..., -2.9617e-01,\n",
       "            4.2663e-01,  1.6154e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.1858e+00, -1.2809e+00, -4.3294e-01,  ..., -3.8595e-01,\n",
       "           -9.2328e-01,  1.6585e-01],\n",
       "          [ 8.2947e-01,  6.3683e-01, -5.3624e-01,  ..., -1.3450e+00,\n",
       "            3.6364e-01, -3.0807e-01],\n",
       "          [-1.5448e-01,  9.8334e-01,  1.2846e-01,  ...,  2.5470e-01,\n",
       "           -6.2691e-01, -1.9639e+00],\n",
       "          ...,\n",
       "          [ 2.8078e-01,  5.2874e-01,  5.8723e-01,  ..., -3.0012e-01,\n",
       "            1.1098e+00, -1.0701e+00],\n",
       "          [ 1.0423e+00, -5.8455e-01, -2.8637e-01,  ...,  7.7986e-01,\n",
       "            3.1214e-01, -1.2657e-02],\n",
       "          [-9.2698e-01,  1.1366e+00, -1.0415e+00,  ...,  1.1512e-02,\n",
       "            1.9024e-01,  5.9322e-01]],\n",
       "\n",
       "         [[ 1.6360e+00, -6.4759e-01, -2.0098e-01,  ..., -1.2945e+00,\n",
       "           -5.9824e-02, -1.0538e+00],\n",
       "          [-1.7403e+00, -4.2303e-01, -1.4791e+00,  ..., -3.7413e-01,\n",
       "            2.8537e-01, -1.2946e+00],\n",
       "          [-8.5910e-01,  1.1556e-01, -7.2728e-01,  ..., -4.9952e-01,\n",
       "           -1.5480e+00,  4.6670e-01],\n",
       "          ...,\n",
       "          [-4.8541e-01, -7.3483e-01, -3.9687e-01,  ..., -3.9528e-01,\n",
       "           -3.0813e-01, -8.6045e-01],\n",
       "          [-4.9138e-01, -1.2478e+00, -1.5475e+00,  ...,  5.6124e-01,\n",
       "            4.3685e-01, -1.2515e+00],\n",
       "          [ 5.2597e-01,  9.3130e-01,  2.1362e-01,  ..., -1.4238e+00,\n",
       "            7.8725e-01, -3.5288e-01]],\n",
       "\n",
       "         [[ 1.2361e+00,  6.4153e-01,  1.5813e+00,  ..., -3.7752e-01,\n",
       "           -5.9345e-01, -4.9646e-01],\n",
       "          [ 9.4146e-01, -1.8114e+00, -4.2762e-01,  ..., -2.4209e+00,\n",
       "            1.6585e-01, -3.4077e-01],\n",
       "          [-7.3200e-01,  1.2912e+00, -1.6085e+00,  ..., -2.2393e+00,\n",
       "           -5.3902e-01,  9.4430e-01],\n",
       "          ...,\n",
       "          [ 2.2207e+00,  8.3638e-01, -6.2100e-01,  ...,  4.9244e-01,\n",
       "            7.7228e-01, -1.9842e+00],\n",
       "          [ 5.4734e-03, -4.5859e-01,  8.2480e-02,  ...,  3.1014e+00,\n",
       "            1.6745e-01, -3.6810e-01],\n",
       "          [-1.4222e+00, -2.1622e+00,  5.2826e-01,  ...,  2.8942e-01,\n",
       "           -3.2123e-01,  1.7621e-01]]],\n",
       "\n",
       "\n",
       "        [[[ 2.4929e-01, -7.4105e-01, -3.8528e-01,  ..., -5.3704e-01,\n",
       "           -1.2575e+00, -2.5695e-01],\n",
       "          [ 1.1686e+00,  4.8796e-01, -1.6944e+00,  ...,  1.1219e+00,\n",
       "           -1.1295e-02, -9.2166e-01],\n",
       "          [ 2.5007e+00, -1.6050e+00,  9.4897e-01,  ..., -3.2944e-01,\n",
       "            8.0536e-01,  7.4953e-01],\n",
       "          ...,\n",
       "          [ 1.3766e-01, -6.7250e-01, -1.5489e+00,  ...,  5.3095e-01,\n",
       "            2.9255e-01,  2.4931e+00],\n",
       "          [-8.1218e-01,  3.4864e-01,  7.3934e-01,  ...,  2.3239e-01,\n",
       "           -2.2687e-01, -1.4470e+00],\n",
       "          [ 5.3062e-01, -1.7223e-01, -1.6042e+00,  ..., -2.6121e+00,\n",
       "            1.4697e-01, -1.4853e+00]],\n",
       "\n",
       "         [[-5.8871e-01, -1.8357e+00, -1.3587e-01,  ..., -7.2910e-03,\n",
       "            1.0069e+00, -1.4972e+00],\n",
       "          [-6.9520e-01,  2.0162e+00,  5.2162e-01,  ...,  2.8667e+00,\n",
       "            1.1102e+00, -7.1531e-01],\n",
       "          [ 2.0704e+00,  2.0079e+00,  4.9047e-01,  ..., -6.2691e-01,\n",
       "            4.6332e-01,  7.3139e-02],\n",
       "          ...,\n",
       "          [ 1.8924e+00, -2.2810e+00,  1.1952e+00,  ..., -3.1790e-01,\n",
       "           -1.4459e+00,  1.2883e+00],\n",
       "          [-7.5157e-01,  1.4611e+00,  4.4859e-01,  ..., -1.5459e+00,\n",
       "           -1.7211e-01, -1.0359e+00],\n",
       "          [ 7.0063e-01, -3.9072e-01,  2.5149e-01,  ...,  1.4935e-01,\n",
       "           -1.0536e+00, -1.4015e+00]],\n",
       "\n",
       "         [[-1.2825e+00, -1.4566e-01, -1.0488e+00,  ...,  6.2289e-01,\n",
       "           -1.7138e-01, -1.0179e+00],\n",
       "          [ 1.6172e+00, -5.2188e-01, -1.2129e-01,  ...,  1.3340e+00,\n",
       "           -1.2119e+00,  8.0458e-01],\n",
       "          [-1.7513e-02,  6.1134e-01, -6.3754e-02,  ...,  4.4689e-01,\n",
       "           -9.0527e-01, -5.3393e-01],\n",
       "          ...,\n",
       "          [ 1.4472e+00,  1.5995e-01, -8.6227e-01,  ..., -4.7212e-01,\n",
       "           -2.5965e-01,  2.3539e+00],\n",
       "          [-4.5403e-01,  9.4185e-02,  9.3265e-01,  ..., -3.8800e-01,\n",
       "           -3.9842e-01, -8.5993e-01],\n",
       "          [ 5.3805e-01,  2.6092e-01, -7.1228e-01,  ..., -2.6376e-01,\n",
       "            2.2993e-01, -3.3059e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-1.1962e+00, -1.3522e-01,  1.5208e-01,  ...,  5.0522e-01,\n",
       "            1.4488e+00, -3.8357e-01],\n",
       "          [-1.3596e+00,  2.2119e+00,  4.8911e-01,  ...,  3.6406e-02,\n",
       "            6.9506e-01, -1.1226e+00],\n",
       "          [-1.0855e+00,  5.1707e-01, -6.9111e-01,  ..., -1.7874e+00,\n",
       "            4.4220e-01,  3.4860e-01],\n",
       "          ...,\n",
       "          [ 1.3761e-02, -1.2795e+00,  3.8998e-03,  ..., -5.5806e-01,\n",
       "           -5.2629e-01, -1.1920e-01],\n",
       "          [-1.7003e+00, -9.5399e-02, -1.4493e+00,  ..., -3.5780e-01,\n",
       "           -5.2854e-01,  4.7687e-01],\n",
       "          [-4.4833e-01, -1.3777e-01,  8.8489e-01,  ..., -7.5936e-02,\n",
       "           -5.4574e-01, -3.1436e-01]],\n",
       "\n",
       "         [[ 2.0018e+00, -6.2949e-01, -1.3551e-01,  ...,  3.4653e-01,\n",
       "           -9.5930e-01,  2.0810e+00],\n",
       "          [-9.9289e-01,  1.4563e+00, -2.3477e+00,  ...,  5.0566e-01,\n",
       "           -2.6746e-01,  7.2568e-01],\n",
       "          [ 2.9169e-01, -1.4978e+00,  2.0113e+00,  ...,  4.7137e-01,\n",
       "           -1.0014e+00,  1.2792e+00],\n",
       "          ...,\n",
       "          [ 1.0726e+00,  3.2483e-01, -1.3123e-01,  ..., -4.5787e-01,\n",
       "            1.6772e+00,  5.4739e-01],\n",
       "          [ 1.9208e-01,  4.0499e-01,  7.9709e-01,  ...,  8.5000e-02,\n",
       "           -9.7822e-02, -7.0994e-01],\n",
       "          [-1.4233e+00, -4.5329e-01,  2.8850e-01,  ..., -2.7925e-01,\n",
       "            3.0177e-01, -1.2955e-01]],\n",
       "\n",
       "         [[-3.1003e-01, -7.4228e-01, -9.5773e-01,  ..., -1.2105e+00,\n",
       "            5.1190e-01,  6.2785e-01],\n",
       "          [-9.0197e-01,  1.6819e+00, -1.5201e+00,  ..., -3.8787e-01,\n",
       "            1.4800e-01, -5.0125e-01],\n",
       "          [-6.8089e-01,  1.8514e-01,  9.2755e-01,  ...,  7.0776e-01,\n",
       "           -3.8985e-01,  5.6512e-01],\n",
       "          ...,\n",
       "          [-6.3951e-01, -1.5425e+00, -7.2424e-01,  ...,  1.4929e+00,\n",
       "            1.2135e+00, -1.6759e+00],\n",
       "          [-1.0960e+00,  3.0551e-01, -1.2938e+00,  ...,  2.2176e-01,\n",
       "           -4.0481e-01,  1.0218e-01],\n",
       "          [-7.9104e-01,  1.8170e+00, -9.8155e-01,  ...,  5.0842e-01,\n",
       "           -2.1633e+00, -1.0655e+00]]],\n",
       "\n",
       "\n",
       "        [[[-5.1267e-01, -6.9006e-02, -1.0766e-01,  ...,  6.8279e-01,\n",
       "            3.4332e+00,  2.3968e+00],\n",
       "          [ 5.0404e-01, -1.6763e+00,  1.5124e+00,  ..., -1.1216e+00,\n",
       "            6.6115e-01, -6.8900e-01],\n",
       "          [ 1.0238e+00, -2.4566e-02,  5.4880e-01,  ..., -8.2523e-03,\n",
       "            1.1191e+00,  8.5146e-01],\n",
       "          ...,\n",
       "          [ 1.9462e+00,  3.6648e-02, -1.0374e+00,  ...,  1.9346e+00,\n",
       "            3.1813e-01, -6.8275e-01],\n",
       "          [-5.9179e-01,  3.6225e-01,  1.8115e-01,  ...,  2.1024e+00,\n",
       "           -1.9599e+00, -1.6031e+00],\n",
       "          [-9.4382e-01,  2.7629e+00, -7.5467e-01,  ...,  2.7714e-02,\n",
       "           -1.0838e+00,  1.3784e+00]],\n",
       "\n",
       "         [[-8.0606e-01, -2.8709e-01, -1.5482e+00,  ...,  1.8709e-01,\n",
       "            2.2496e+00,  2.4322e-01],\n",
       "          [-1.8889e-01,  9.0891e-01,  5.8307e-01,  ..., -5.9215e-01,\n",
       "           -8.0817e-01, -3.8920e-01],\n",
       "          [-8.1550e-01, -2.0237e+00, -3.3292e-01,  ..., -4.9662e-01,\n",
       "            9.0356e-01, -1.3665e-01],\n",
       "          ...,\n",
       "          [-1.9742e-01, -1.0669e+00,  2.0383e-01,  ..., -1.6713e+00,\n",
       "            1.8541e+00,  5.2016e-01],\n",
       "          [ 3.8764e-01, -2.4980e-01, -2.0525e+00,  ..., -2.7393e-01,\n",
       "            5.8920e-01, -9.4769e-01],\n",
       "          [ 2.2365e-01,  1.1743e+00,  2.9540e-01,  ...,  4.0334e-01,\n",
       "           -2.5302e-02, -2.2546e-01]],\n",
       "\n",
       "         [[-9.9103e-01, -2.5803e-01,  7.9875e-01,  ..., -1.1397e+00,\n",
       "           -6.3262e-01,  7.9848e-01],\n",
       "          [-1.8061e+00,  3.2324e-01, -3.5662e-01,  ..., -1.0106e+00,\n",
       "           -4.2172e-01,  7.1171e-01],\n",
       "          [ 4.3481e-01,  1.9421e+00, -1.8376e-01,  ...,  3.5263e-01,\n",
       "            9.8578e-02, -1.0686e+00],\n",
       "          ...,\n",
       "          [-6.5127e-01,  5.6866e-01,  5.3052e-02,  ...,  1.5188e-01,\n",
       "           -1.0681e-01, -1.2761e+00],\n",
       "          [-1.4842e+00, -3.1089e-01, -2.1545e-01,  ...,  2.1145e+00,\n",
       "           -1.4519e+00, -9.7791e-01],\n",
       "          [-1.8920e+00,  3.7313e-01, -8.4075e-01,  ...,  1.5458e-01,\n",
       "            8.4641e-02, -5.7135e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 1.7265e+00,  4.5745e-01,  3.5000e-01,  ...,  3.1505e-03,\n",
       "           -1.0047e+00,  1.2153e+00],\n",
       "          [ 1.0790e+00, -6.3171e-01, -2.8225e-01,  ..., -8.0379e-02,\n",
       "            1.0075e+00,  1.3921e+00],\n",
       "          [ 1.3312e+00, -1.5865e+00, -2.4552e-01,  ...,  1.4781e+00,\n",
       "            7.9633e-01,  5.2642e-01],\n",
       "          ...,\n",
       "          [-2.4473e+00, -7.0648e-01, -1.9083e-01,  ...,  1.1311e+00,\n",
       "            1.4231e+00, -7.2126e-01],\n",
       "          [ 1.3362e-01,  1.4043e+00,  6.0891e-03,  ...,  1.8031e+00,\n",
       "           -9.0119e-02, -8.6487e-01],\n",
       "          [ 1.9248e+00,  1.2260e+00,  1.5756e+00,  ..., -1.9314e+00,\n",
       "           -1.5983e-01,  9.7294e-02]],\n",
       "\n",
       "         [[-2.9018e+00,  1.2688e-01, -1.7567e-01,  ..., -1.9938e-01,\n",
       "            4.2152e-01,  1.1861e+00],\n",
       "          [ 1.6539e-01, -1.7820e+00,  1.6490e-01,  ...,  1.2641e+00,\n",
       "            1.0225e+00, -7.1815e-01],\n",
       "          [-4.6790e-02, -4.4126e-01,  2.5669e-01,  ..., -5.1352e-01,\n",
       "           -3.4507e-01,  1.1831e+00],\n",
       "          ...,\n",
       "          [ 2.7046e-01, -1.8769e-01, -6.0478e-01,  ..., -1.3872e-01,\n",
       "            8.3271e-02, -1.3749e-01],\n",
       "          [ 5.9055e-02,  1.8833e-01, -3.1737e-01,  ...,  4.3277e-01,\n",
       "            7.6034e-01, -2.8370e-01],\n",
       "          [-4.8149e-01,  3.4277e-01, -1.7723e-01,  ..., -9.0403e-01,\n",
       "            7.5498e-01,  4.1175e-01]],\n",
       "\n",
       "         [[ 2.3381e-01,  6.6450e-01, -4.6151e-01,  ...,  1.4064e+00,\n",
       "           -5.3340e-01,  1.3827e+00],\n",
       "          [ 2.9628e+00,  1.7551e+00,  1.8675e+00,  ...,  9.8440e-03,\n",
       "            3.4014e-01,  8.6654e-01],\n",
       "          [-4.4247e-01, -1.3447e+00, -5.3637e-01,  ..., -5.5113e-01,\n",
       "            1.4419e+00, -6.5676e-01],\n",
       "          ...,\n",
       "          [ 3.3439e-01,  9.6413e-03,  1.4359e+00,  ...,  6.4906e-01,\n",
       "            5.7021e-01,  1.9068e-01],\n",
       "          [-3.4851e-01,  8.6501e-01,  2.4914e-01,  ...,  6.5022e-01,\n",
       "            1.0644e+00, -1.2050e+00],\n",
       "          [-1.4408e+00,  9.4563e-01,  1.4338e+00,  ..., -1.8618e-01,\n",
       "            1.3361e+00, -1.5787e-01]]]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.reshape(3,501,32,129)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 1.4173e+00,  7.7962e-01,  3.9546e-03,  ..., -2.7568e-02,\n",
       "           -9.6422e-01, -3.3995e-01],\n",
       "          [-4.1776e-01,  2.0515e-01, -3.2592e-01,  ..., -5.0554e-01,\n",
       "           -8.5626e-01, -7.5003e-01],\n",
       "          [-8.8587e-01,  1.3755e+00,  1.5670e+00,  ..., -6.8439e-01,\n",
       "           -8.3740e-01,  3.2749e-01],\n",
       "          ...,\n",
       "          [-1.4022e+00,  1.0135e+00,  1.8463e-01,  ..., -1.2581e+00,\n",
       "           -2.2503e-01, -2.0165e+00],\n",
       "          [ 2.5633e-01, -2.2823e-01, -1.7896e+00,  ...,  1.2296e+00,\n",
       "            2.0197e-01,  1.6196e+00],\n",
       "          [ 2.0377e+00,  2.0726e-02, -5.2761e-01,  ..., -8.6264e-01,\n",
       "            3.5127e-01,  1.1185e+00]],\n",
       "\n",
       "         [[ 1.4753e-02, -1.5288e+00,  5.1300e-01,  ..., -1.9317e+00,\n",
       "            2.2782e+00, -2.8368e-01],\n",
       "          [-8.7368e-01, -7.2315e-01,  1.3629e+00,  ...,  1.0970e-01,\n",
       "           -1.7061e+00,  4.5335e-01],\n",
       "          [-6.9197e-02, -1.7097e+00,  9.4216e-01,  ...,  1.3155e+00,\n",
       "           -1.0831e+00, -1.8865e+00],\n",
       "          ...,\n",
       "          [ 1.7914e-02,  2.8952e-01,  8.8451e-01,  ..., -1.0918e+00,\n",
       "            1.3390e+00, -1.3086e-01],\n",
       "          [-8.5710e-01, -6.0583e-01,  2.7612e-01,  ...,  8.9928e-02,\n",
       "            1.1961e+00, -2.8358e-01],\n",
       "          [ 1.1638e+00,  3.5013e-01,  5.3242e-01,  ...,  1.4750e+00,\n",
       "           -2.7274e+00,  3.0799e+00]],\n",
       "\n",
       "         [[ 1.6739e+00,  1.0315e+00, -3.1527e-01,  ..., -1.3023e+00,\n",
       "           -1.5804e+00, -3.4867e-01],\n",
       "          [ 9.8544e-01, -5.2546e-01,  1.9223e+00,  ..., -6.8322e-01,\n",
       "           -5.6756e-01,  7.2929e-01],\n",
       "          [-2.2090e+00,  1.1775e+00,  8.7367e-01,  ..., -2.6072e-01,\n",
       "            5.9265e-01, -7.8509e-02],\n",
       "          ...,\n",
       "          [ 3.6346e-01, -7.2150e-01, -4.5104e-01,  ..., -2.3210e-01,\n",
       "            1.1761e+00,  9.5823e-01],\n",
       "          [ 5.0664e-01, -2.0768e-02,  9.3340e-01,  ...,  8.9009e-01,\n",
       "           -1.0256e+00,  2.5602e+00],\n",
       "          [-2.8030e-01, -4.0707e-02, -4.3855e-01,  ...,  1.4687e+00,\n",
       "            7.8256e-01,  8.8403e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 3.8601e-01,  2.8655e-01,  2.1920e+00,  ...,  7.1205e-01,\n",
       "            2.7848e-01, -8.3867e-01],\n",
       "          [ 5.9739e-01, -9.2320e-01,  1.5767e+00,  ..., -6.5097e-01,\n",
       "            2.1796e-01, -1.0473e+00],\n",
       "          [ 1.3937e+00,  2.2259e-01, -1.1127e+00,  ..., -2.0836e-01,\n",
       "           -9.9526e-02,  6.5295e-01],\n",
       "          ...,\n",
       "          [ 2.0137e+00,  1.5381e+00, -6.5393e-01,  ..., -3.1760e-01,\n",
       "            4.1427e-01, -5.9424e-01],\n",
       "          [-2.9314e-01, -9.8124e-02,  7.1954e-01,  ..., -1.3601e-01,\n",
       "            4.9403e-01, -1.1352e+00],\n",
       "          [-1.1981e-01,  1.8555e+00, -1.5615e+00,  ...,  9.3173e-01,\n",
       "           -3.7292e-01, -2.4580e-01]],\n",
       "\n",
       "         [[ 1.8548e+00,  3.8566e-02,  5.5926e-01,  ..., -7.4575e-01,\n",
       "            3.3124e-01,  1.0622e+00],\n",
       "          [-6.7331e-01, -2.5818e-01,  5.1046e-01,  ..., -1.0622e+00,\n",
       "           -4.0790e-01,  1.4873e-01],\n",
       "          [ 2.5819e-01,  5.2128e-01,  3.2406e-01,  ..., -1.1059e+00,\n",
       "           -1.1847e+00,  7.3285e-01],\n",
       "          ...,\n",
       "          [-1.5664e+00,  1.8141e-01,  1.7125e+00,  ...,  4.0256e-01,\n",
       "            3.8353e-01, -1.2825e-02],\n",
       "          [-1.8569e+00,  3.3713e+00, -3.9484e-01,  ..., -5.0453e-01,\n",
       "            5.5561e-01, -8.8745e-01],\n",
       "          [ 7.3779e-01, -4.7684e-01,  8.6579e-01,  ...,  9.2439e-01,\n",
       "            1.0300e+00,  5.7792e-01]],\n",
       "\n",
       "         [[ 2.6070e-01, -2.7698e+00,  1.0541e+00,  ..., -1.7635e+00,\n",
       "            2.7667e-01, -8.1514e-01],\n",
       "          [-2.9682e-01,  1.7731e+00,  8.8659e-01,  ..., -1.4869e-01,\n",
       "            9.5589e-01, -4.3438e-01],\n",
       "          [ 1.8267e+00, -8.5254e-01,  1.4114e+00,  ...,  9.4400e-01,\n",
       "            9.3839e-01,  3.7018e-01],\n",
       "          ...,\n",
       "          [ 2.2207e+00,  8.3638e-01, -6.2100e-01,  ...,  4.9244e-01,\n",
       "            7.7228e-01, -1.9842e+00],\n",
       "          [ 5.4734e-03, -4.5859e-01,  8.2480e-02,  ...,  3.1014e+00,\n",
       "            1.6745e-01, -3.6810e-01],\n",
       "          [-1.4222e+00, -2.1622e+00,  5.2826e-01,  ...,  2.8942e-01,\n",
       "           -3.2123e-01,  1.7621e-01]]],\n",
       "\n",
       "\n",
       "        [[[ 2.4929e-01, -7.4105e-01, -3.8528e-01,  ..., -5.3704e-01,\n",
       "           -1.2575e+00, -2.5695e-01],\n",
       "          [ 1.1686e+00,  4.8796e-01, -1.6944e+00,  ...,  1.1219e+00,\n",
       "           -1.1295e-02, -9.2166e-01],\n",
       "          [ 2.5007e+00, -1.6050e+00,  9.4897e-01,  ..., -3.2944e-01,\n",
       "            8.0536e-01,  7.4953e-01],\n",
       "          ...,\n",
       "          [-4.0642e-01, -1.1269e+00, -2.3018e-01,  ..., -7.9425e-01,\n",
       "            7.0871e-01,  1.3945e+00],\n",
       "          [-8.5102e-01,  1.7688e-01, -6.9886e-01,  ..., -8.9247e-01,\n",
       "            4.3439e-01, -3.3122e-01],\n",
       "          [ 2.9158e-01,  9.1199e-01, -1.1841e+00,  ..., -6.1248e-01,\n",
       "           -1.4973e+00, -8.3326e-01]],\n",
       "\n",
       "         [[-1.4022e+00,  3.2680e-01,  2.2337e-01,  ...,  1.4585e+00,\n",
       "            6.4998e-01,  1.7069e+00],\n",
       "          [-9.2758e-01,  2.2306e-01,  4.8614e-01,  ..., -4.8598e-01,\n",
       "           -1.4673e+00, -9.4140e-01],\n",
       "          [ 6.3466e-01,  1.2991e+00,  1.2315e+00,  ..., -2.4192e-01,\n",
       "           -6.2871e-01, -5.2704e-01],\n",
       "          ...,\n",
       "          [ 1.3431e+00, -7.3320e-01, -2.6601e-02,  ..., -1.0094e-01,\n",
       "           -1.2005e+00,  4.6383e-01],\n",
       "          [-1.1769e+00,  3.0025e-01,  1.2437e+00,  ...,  1.7800e+00,\n",
       "           -8.3851e-01,  1.3130e+00],\n",
       "          [ 2.4963e-01,  1.6413e+00, -1.0033e+00,  ...,  2.0138e+00,\n",
       "           -8.0449e-02, -1.1304e+00]],\n",
       "\n",
       "         [[ 3.4997e-01,  1.9013e-01, -7.6698e-01,  ..., -1.1074e+00,\n",
       "            6.7564e-02, -5.9383e-01],\n",
       "          [-5.6101e-01, -1.3368e+00,  5.4323e-01,  ..., -7.1809e-01,\n",
       "           -6.4228e-01,  5.8490e-02],\n",
       "          [-8.1075e-01, -1.6568e+00, -1.3387e+00,  ..., -1.5253e+00,\n",
       "            1.1223e+00,  1.6390e+00],\n",
       "          ...,\n",
       "          [ 2.0412e+00,  7.8006e-01, -1.0693e-01,  ...,  1.2924e+00,\n",
       "           -2.1654e+00, -8.7400e-01],\n",
       "          [-1.7197e+00,  9.2442e-01, -2.7010e-02,  ..., -1.0712e+00,\n",
       "           -5.4550e-01, -3.9894e-01],\n",
       "          [-8.1815e-01, -8.3669e-01, -1.5639e-01,  ...,  9.4761e-01,\n",
       "            4.8254e-01,  3.9436e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-3.2489e-02,  1.8905e-02,  2.1748e-01,  ..., -6.9110e-02,\n",
       "           -5.2196e-01, -1.1462e+00],\n",
       "          [ 4.7192e-01,  5.4301e-01,  2.4795e+00,  ..., -1.8399e-01,\n",
       "           -6.5124e-02, -9.5010e-01],\n",
       "          [-2.8719e-02,  4.3024e-01,  9.3862e-01,  ...,  5.4213e-01,\n",
       "           -7.3249e-01, -1.5958e+00],\n",
       "          ...,\n",
       "          [-1.0279e+00, -1.0983e+00,  8.7479e-01,  ...,  2.5433e-01,\n",
       "            3.0129e-01, -7.3457e-01],\n",
       "          [-1.5610e-01, -1.6793e+00, -1.7649e+00,  ..., -2.7919e-01,\n",
       "            1.9859e+00,  3.8414e-01],\n",
       "          [-1.5601e+00, -2.1508e-01,  9.1194e-01,  ...,  1.2076e+00,\n",
       "           -8.1575e-01,  2.0569e+00]],\n",
       "\n",
       "         [[ 5.4541e-01, -4.6309e-01,  9.1975e-01,  ...,  3.2859e-01,\n",
       "           -1.7595e+00,  5.8060e-01],\n",
       "          [ 4.9083e-01,  1.1366e+00,  5.8537e-01,  ..., -1.3149e-01,\n",
       "            6.2809e-01, -1.6355e+00],\n",
       "          [-2.5835e-01, -3.1178e-01, -1.7765e+00,  ..., -9.5343e-02,\n",
       "            1.4687e+00,  1.9439e+00],\n",
       "          ...,\n",
       "          [-4.3066e-01,  1.4131e+00, -1.2286e-01,  ..., -1.1717e+00,\n",
       "            1.8165e-01,  2.3576e+00],\n",
       "          [-1.4112e+00,  1.0448e-01, -6.9705e-01,  ..., -2.0637e-01,\n",
       "            7.4937e-01, -3.7116e-01],\n",
       "          [-3.9438e-01,  7.9560e-01,  1.1906e+00,  ..., -3.4532e-01,\n",
       "           -1.1951e-01,  2.0174e+00]],\n",
       "\n",
       "         [[ 2.6423e-01,  8.4896e-01, -6.4185e-01,  ..., -9.0316e-01,\n",
       "            4.2255e-01,  1.1125e+00],\n",
       "          [ 3.1544e-01, -1.7384e-02,  1.2149e+00,  ...,  1.0241e+00,\n",
       "            6.0655e-01,  3.1697e-01],\n",
       "          [ 7.5223e-01, -1.5645e+00,  7.4356e-01,  ...,  2.1007e+00,\n",
       "            1.2704e+00,  5.2666e-01],\n",
       "          ...,\n",
       "          [-6.3951e-01, -1.5425e+00, -7.2424e-01,  ...,  1.4929e+00,\n",
       "            1.2135e+00, -1.6759e+00],\n",
       "          [-1.0960e+00,  3.0551e-01, -1.2938e+00,  ...,  2.2176e-01,\n",
       "           -4.0481e-01,  1.0218e-01],\n",
       "          [-7.9104e-01,  1.8170e+00, -9.8155e-01,  ...,  5.0842e-01,\n",
       "           -2.1633e+00, -1.0655e+00]]],\n",
       "\n",
       "\n",
       "        [[[-5.1267e-01, -6.9006e-02, -1.0766e-01,  ...,  6.8279e-01,\n",
       "            3.4332e+00,  2.3968e+00],\n",
       "          [ 5.0404e-01, -1.6763e+00,  1.5124e+00,  ..., -1.1216e+00,\n",
       "            6.6115e-01, -6.8900e-01],\n",
       "          [ 1.0238e+00, -2.4566e-02,  5.4880e-01,  ..., -8.2523e-03,\n",
       "            1.1191e+00,  8.5146e-01],\n",
       "          ...,\n",
       "          [-1.2259e+00, -2.1323e-01,  5.9049e-01,  ..., -5.2353e-01,\n",
       "            1.0952e+00,  4.8396e-02],\n",
       "          [ 2.7138e-01, -4.5839e-01,  3.6397e-01,  ...,  8.8268e-02,\n",
       "            1.5202e+00,  6.8128e-01],\n",
       "          [ 7.1377e-01, -6.3030e-01,  3.0845e-01,  ..., -3.0431e+00,\n",
       "           -1.2878e-01,  3.7702e-01]],\n",
       "\n",
       "         [[ 1.3211e+00, -3.0770e-01, -2.8367e-01,  ..., -1.1510e+00,\n",
       "           -8.2442e-01,  9.0989e-01],\n",
       "          [-1.8303e-01,  1.2330e+00, -8.4881e-01,  ...,  4.9534e-01,\n",
       "           -5.2415e-01, -1.3119e+00],\n",
       "          [ 7.8133e-01,  3.7116e-01,  6.7050e-01,  ...,  4.6082e-01,\n",
       "           -1.2753e+00, -3.7845e-02],\n",
       "          ...,\n",
       "          [-2.6448e-01, -3.6857e-01,  6.2289e-01,  ...,  3.0307e-01,\n",
       "            5.7001e-01,  3.1135e-01],\n",
       "          [ 1.8374e-01, -7.7393e-01,  5.3559e-02,  ..., -3.9925e-02,\n",
       "           -2.5661e-02, -4.3909e-01],\n",
       "          [ 6.6308e-01,  5.4135e-01,  7.3519e-01,  ..., -1.0981e+00,\n",
       "           -1.0813e-01, -1.1063e+00]],\n",
       "\n",
       "         [[-5.9182e-01, -2.0487e-01, -2.3376e-01,  ...,  3.7870e-01,\n",
       "           -1.3042e+00,  2.3440e-01],\n",
       "          [ 3.9597e-01,  8.5352e-01,  2.6053e+00,  ..., -5.4151e-01,\n",
       "           -1.5815e-01,  4.4251e-01],\n",
       "          [-3.4286e-03,  4.0033e-01, -2.7187e-01,  ..., -1.5198e+00,\n",
       "           -2.7401e+00,  3.2023e+00],\n",
       "          ...,\n",
       "          [-8.9855e-01, -2.1322e+00,  5.8289e-01,  ...,  4.3034e-01,\n",
       "            4.1946e-01,  6.4346e-01],\n",
       "          [ 4.4918e-01,  2.3160e+00,  2.6619e-01,  ...,  1.7723e+00,\n",
       "            2.4163e-01,  5.2962e-01],\n",
       "          [ 5.9954e-01, -6.0741e-01, -1.2941e+00,  ..., -1.7115e+00,\n",
       "            1.2895e+00, -2.2337e-01]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-1.4712e+00, -8.0756e-01,  7.2554e-01,  ...,  1.7935e+00,\n",
       "           -1.8760e-01, -1.4849e-01],\n",
       "          [ 2.0135e+00,  2.4961e-01, -4.3150e-01,  ..., -1.4594e+00,\n",
       "           -8.1198e-02,  3.0800e+00],\n",
       "          [ 1.2727e+00,  4.7884e-01,  3.5763e-01,  ...,  6.1813e-01,\n",
       "           -1.5325e+00, -2.6912e+00],\n",
       "          ...,\n",
       "          [-1.4847e+00, -6.6559e-01, -1.3845e+00,  ..., -3.9385e-01,\n",
       "            3.5611e-01, -4.2478e-01],\n",
       "          [-1.2053e-01, -4.9656e-01, -9.0624e-01,  ...,  8.3161e-01,\n",
       "            5.8515e-01, -1.0430e+00],\n",
       "          [-5.1855e-01, -4.3050e-01,  4.5714e-01,  ..., -1.9104e-01,\n",
       "            4.9361e-01, -6.9194e-01]],\n",
       "\n",
       "         [[-4.3663e-01, -1.2228e+00, -9.6029e-01,  ..., -7.3643e-01,\n",
       "            4.1701e-02,  4.4251e-01],\n",
       "          [-1.2423e-01,  2.5922e-01,  2.6780e+00,  ..., -8.3903e-01,\n",
       "            4.8566e-01,  9.8491e-01],\n",
       "          [ 1.1225e+00,  6.4786e-01,  4.2604e-02,  ..., -9.1205e-03,\n",
       "           -1.2349e+00,  1.8163e-01],\n",
       "          ...,\n",
       "          [ 4.5860e-01, -3.7439e-01,  5.1245e-01,  ...,  4.9650e-01,\n",
       "           -2.9094e+00, -5.0044e-01],\n",
       "          [ 1.3822e+00,  5.4450e-01, -6.2038e-02,  ...,  2.0740e+00,\n",
       "           -5.7913e-01, -5.2395e-01],\n",
       "          [-1.8508e-01, -1.9053e-01, -1.2883e-01,  ..., -6.1952e-01,\n",
       "            1.9415e+00,  2.9476e-01]],\n",
       "\n",
       "         [[ 1.4926e+00, -9.0899e-01,  9.1773e-01,  ..., -1.9788e+00,\n",
       "            7.5944e-01,  1.3949e-01],\n",
       "          [ 7.5506e-01, -3.8522e-01,  1.0227e-01,  ..., -7.7079e-01,\n",
       "           -7.5733e-01,  6.7085e-01],\n",
       "          [ 3.4589e-01,  3.4268e-01,  8.3856e-01,  ..., -9.1128e-01,\n",
       "           -9.9181e-01, -1.0451e+00],\n",
       "          ...,\n",
       "          [ 3.3439e-01,  9.6413e-03,  1.4359e+00,  ...,  6.4906e-01,\n",
       "            5.7021e-01,  1.9068e-01],\n",
       "          [-3.4851e-01,  8.6501e-01,  2.4914e-01,  ...,  6.5022e-01,\n",
       "            1.0644e+00, -1.2050e+00],\n",
       "          [-1.4408e+00,  9.4563e-01,  1.4338e+00,  ..., -1.8618e-01,\n",
       "            1.3361e+00, -1.5787e-01]]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.reshape(3,32,501,129)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 2, 2, 501, 129])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(3,2,2,501,129)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = torch.complex(x[:,:,0], x[:,:,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 2, 501, 129])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = y.view(3*2, 501, 129)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 501, 129])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loss import loss\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.2318, 0.1650],\n",
      "         [0.7633, 0.8316]],\n",
      "\n",
      "        [[0.6742, 0.8293],\n",
      "         [0.8035, 0.6347]]])\n",
      "tensor([[[0.2243, 0.4870],\n",
      "         [0.3446, 0.1579]],\n",
      "\n",
      "        [[0.4255, 0.5159],\n",
      "         [0.4889, 0.0432]]])\n"
     ]
    }
   ],
   "source": [
    "source = torch.rand(2, 2, 2)\n",
    "estimate_source = torch.rand(2,2,2)\n",
    "print(source)\n",
    "print(estimate_source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-0.4353,  0.4353],\n",
      "          [ 0.4630, -0.4630]]],\n",
      "\n",
      "\n",
      "        [[[-0.0238,  0.0238],\n",
      "          [ 1.0929, -1.0929]]]])\n",
      "tensor([[[[-0.8729,  0.8729]],\n",
      "\n",
      "         [[ 0.9123, -0.9123]]],\n",
      "\n",
      "\n",
      "        [[[-0.8412,  0.8412]],\n",
      "\n",
      "         [[-0.0769,  0.0769]]]])\n",
      "s target torch.Size([2, 1, 2, 1])\n",
      "tensor([[[[3.7900e-01],\n",
      "          [4.2867e-01]]],\n",
      "\n",
      "\n",
      "        [[[1.1343e-03],\n",
      "          [2.3889e+00]]]])\n",
      "alpha shape torch.Size([2, 2, 2, 1])\n",
      "denominator torch.Size([2, 2, 2, 1])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "torch.manual_seed(878373947641000)\n",
    "from itertools import permutations\n",
    "\n",
    "\n",
    "def get_mask(source, source_lengths):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        source: [B, C, T]\n",
    "        source_lengths: [B]\n",
    "    Returns:\n",
    "        mask: [B, 1, T]\n",
    "    \"\"\"\n",
    "    B, _, T = source.size()\n",
    "    mask = source.new_ones((B, 1, T))\n",
    "    for i in range(B):\n",
    "        mask[i, :, source_lengths[i][0] :] = 0\n",
    "    return mask\n",
    "\n",
    "\n",
    "source = torch.randn(2, 2, 2)\n",
    "estimate_source = torch.randn(2, 2, 2)\n",
    "source_lengths = source_lengths = torch.full(\n",
    "    (estimate_source.size(0), 1), estimate_source.size(2)\n",
    ").to(\"cpu\")\n",
    "\n",
    "\n",
    "assert source.size() == estimate_source.size()\n",
    "B, C, T = source.size()\n",
    "# mask padding position along T\n",
    "mask = get_mask(source, source_lengths)\n",
    "estimate_source *= mask\n",
    "\n",
    "# Step 1. Zero-mean norm\n",
    "num_samples = source_lengths.view(-1, 1, 1).float()  # [B, 1, 1]\n",
    "mean_target = torch.sum(source, dim=2, keepdim=True) / num_samples\n",
    "mean_estimate = torch.sum(estimate_source, dim=2, keepdim=True) / num_samples\n",
    "zero_mean_target = source - mean_target\n",
    "zero_mean_estimate = estimate_source - mean_estimate\n",
    "# mask padding position along T\n",
    "zero_mean_target *= mask\n",
    "zero_mean_estimate *= mask\n",
    "\n",
    "# Step 2. SI-SNR with PIT\n",
    "# reshape to use broadcast\n",
    "s_target = torch.unsqueeze(zero_mean_target, dim=1)  # [B, 1, C, T]\n",
    "s_estimate = torch.unsqueeze(zero_mean_estimate, dim=2)  # [B, C, 1, T]\n",
    "print(s_target)\n",
    "print(s_estimate)\n",
    "\n",
    "s_target_square = torch.sum(s_target**2, dim=3, keepdim=True)  # [B,1,C,1]\n",
    "print(f\"s target {s_target_square.shape}\")\n",
    "print(s_target_square)\n",
    "\n",
    "alpha = torch.sum(s_estimate * s_target, dim=3, keepdim=True) / torch.sum(\n",
    "    s_estimate**2, dim=3, keepdim=True\n",
    ")  # [B,C,C,1]\n",
    "print(f\"alpha shape {alpha.shape}\")\n",
    "\n",
    "denominator = alpha * s_estimate - s_target\n",
    "denominator = torch.sum(denominator**2, dim=3, keepdim=True)  # [B,C,C,T]\n",
    "\n",
    "pair_wise_si_snr = s_target_square / denominator\n",
    "\n",
    "pair_wise_si_snr = 10 * torch.log10(pair_wise_si_snr)\n",
    "pair_wise_si_snr = pair_wise_si_snr.squeeze(3)\n",
    "\n",
    "print(f\"denominator {denominator.shape}\")\n",
    "\n",
    "# Get max_snr of each utterance\n",
    "# permutations, [C!, C]\n",
    "perms = source.new_tensor(list(permutations(range(C))), dtype=torch.long)\n",
    "# one-hot, [C!, C, C]\n",
    "index = torch.unsqueeze(perms, 2)\n",
    "perms_one_hot = source.new_zeros((*perms.size(), C)).scatter_(2, index, 1)\n",
    "# [B, C!] <- [B, C, C] einsum [C!, C, C], SI-SNR sum of each permutation\n",
    "snr_set = torch.einsum(\"bij,pij->bp\", [pair_wise_si_snr, perms_one_hot])\n",
    "max_snr_idx = torch.argmax(snr_set, dim=1)  # [B]\n",
    "# max_snr = torch.gather(snr_set, 1, max_snr_idx.view(-1, 1))  # [B, 1]\n",
    "max_snr, _ = torch.max(snr_set, dim=1, keepdim=True)\n",
    "max_snr /= C\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.4987],\n",
       "         [-0.5304]],\n",
       "\n",
       "        [[-0.4771],\n",
       "         [ 0.5074]]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4353,  0.4353])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(alpha * s_estimate)[0][0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.4353,  0.4353],\n",
       "          [ 0.4630, -0.4630]]],\n",
       "\n",
       "\n",
       "        [[[-0.0238,  0.0238],\n",
       "          [ 1.0929, -1.0929]]]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.4353,  0.4353],\n",
       "         [ 0.4630, -0.4630]]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_target[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-2.9802e-08,  0.0000e+00],\n",
       "          [ 8.9828e-01, -8.9828e-01]],\n",
       "\n",
       "         [[-8.9828e-01,  8.9828e-01],\n",
       "          [-2.9802e-08, -2.9802e-08]]],\n",
       "\n",
       "\n",
       "        [[[ 3.7253e-09, -1.8626e-09],\n",
       "          [ 1.1167e+00, -1.1167e+00]],\n",
       "\n",
       "         [[-1.1167e+00,  1.1167e+00],\n",
       "          [-2.3842e-07,  0.0000e+00]]]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(alpha * s_estimate - s_target.transpose(1,2))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
